﻿
//////////////////////////////////////////////////////////////////////////////
//
// Copyright © 1998-2024 Glodon Company Limited.
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the “Software”),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
// sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
//
//////////////////////////////////////////////////////////////////////////////
#include "LogManager.h"

#include <QMap>
#include <sstream>
#include <QMessageBox>
#include "CrashHelperEx.h"

#pragma warning(push)
#pragma warning(disable: 4091)
#include <DbgHelp.h>
#pragma warning(pop)

#include "EnableCompileWarning_The_LAST_IncludeInCpp.h"

#pragma comment(lib, "dbghelp.lib") 

struct AddressCacheValue
{
    AddressCacheValue(size_t size, DWORD protect)
        : Size(size)
        , Protect(protect)
    {}

    size_t Size;
    DWORD Protect;
};

static LPTOP_LEVEL_EXCEPTION_FILTER OldExceptionFilter = nullptr;
static QMap<void*, AddressCacheValue> AddressCache;

bool CrashHelperEx::IsValidAddress(const void* ptr)
{
    MEMORY_BASIC_INFORMATION mbi = {0};
    if (::VirtualQuery(ptr, &mbi, sizeof(mbi)))
    {
        DWORD mask = (PAGE_READONLY | PAGE_READWRITE | PAGE_WRITECOPY | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY);
        bool b = (mbi.Protect & mask) > 0;
        if (mbi.Protect & (PAGE_GUARD | PAGE_NOACCESS)) 
            b = false;

        return b;
    }

    return false;
}

void CrashHelperEx::CrashWarn(int expression, const QString& message)
{
    if (!expression)
    {
        QMessageBox::critical(nullptr, QString("Crash"), message);
        CrashHelperEx::LogCallStack();
    }
}

void CrashHelperEx::ProtectAddress(void* ptr, size_t size)
{
    DWORD oldProtect;
    DWORD newProtect = PAGE_NOACCESS;
    if (::VirtualProtect(ptr, size, newProtect, &oldProtect))
    {
        AddressCache.insert(ptr, AddressCacheValue(size, oldProtect));
    }
}

LONG WINAPI UnhandledCrashFilter(LPEXCEPTION_POINTERS pException)
{
    if (pException)
    {
        EXCEPTION_RECORD* pRecord = pException->ExceptionRecord;
        if (pRecord)
        {
            if (pRecord->ExceptionCode == EXCEPTION_ACCESS_VIOLATION && pRecord->NumberParameters >= 2)
            {
                ULONG_PTR address = pRecord->ExceptionInformation[1];
                void* ptr = reinterpret_cast<void*>(address);
                if (ptr)
                {
                    auto iFind = AddressCache.find(ptr);
                    if (iFind != AddressCache.end())
                    {
                        size_t size = iFind.value().Size;
                        DWORD oldProtect = iFind.value().Protect;
                        DWORD newProtect;
                        if (::VirtualProtect(ptr, size, oldProtect, &newProtect))
                        {
                            return EXCEPTION_CONTINUE_EXECUTION;
                        }
                    }
                }
            }
        }
    }

    return OldExceptionFilter ? OldExceptionFilter(pException) : EXCEPTION_CONTINUE_SEARCH;
}

void CrashHelperEx::InstallUnhandlExceptionFilter()
{
    OldExceptionFilter = ::SetUnhandledExceptionFilter(UnhandledCrashFilter);
}

void CrashHelperEx::LogCallStack()
{
    HANDLE hProcess = GetCurrentProcess();

    CONTEXT context;
    context.ContextFlags = CONTEXT_FULL;
    RtlCaptureContext(&context);

    SymInitialize(hProcess, NULL, TRUE);
    SymSetOptions(SYMOPT_LOAD_LINES);

    STACKFRAME64 stackFrame;
    ZeroMemory(&stackFrame, sizeof(STACKFRAME64));
    stackFrame.AddrPC.Offset = context.Rip;
    stackFrame.AddrPC.Mode = AddrModeFlat;
    stackFrame.AddrFrame.Offset = context.Rbp;
    stackFrame.AddrFrame.Mode = AddrModeFlat;
    stackFrame.AddrStack.Offset = context.Rsp;
    stackFrame.AddrStack.Mode = AddrModeFlat;

    FYI(L"[Call Stack:]");
    
    HANDLE hThread = GetCurrentThread();
    while (::StackWalk64(IMAGE_FILE_MACHINE_AMD64, hProcess, hThread, &stackFrame, &context, NULL, 0, 0, NULL))
    {
        DWORD64 retAddress = stackFrame.AddrPC.Offset;
        HMODULE hModule;
        ::GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, reinterpret_cast<LPCWSTR>(retAddress), &hModule);
        wchar_t moduleFileName[MAX_PATH] = { 0 };
        GetModuleFileNameW(hModule, moduleFileName, MAX_PATH);
        std::wstringstream wss;
        wss.setf(std::ios::hex, std::ios::basefield);
        wss.width(16);
        wss.fill(L'0');
        wss << retAddress;
        std::wstringstream outs;
        outs << L"[" << moduleFileName << "!0x" << wss.str() << "]";
        FYI(outs.str());
    }
}

